# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jshin

import logging

from coin.base.datetime_util import to_datetime
from coin.base.timestamp import get_timestamp
from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.hk_hard_arb.dump_unhedge_mm import FeatureRow
from coin.strategy.mm.calculator import MovingSum


class EdgeCalculator(object):
  def __init__(self, arb_product, *, price_det_notional):
    self._logger = logging.getLogger(__name__)
    self._arb_product = arb_product
    self._price_det_notional = price_det_notional

    self._books = [None, None]

    self.edge_1_0 = None
    self.edge_0_1 = None
    self.ask_exchange0 = None
    self.bid_exchange0 = None
    self.ask_exchange1 = None
    self.bid_exchange1 = None
    self.ready = False

  def on_book(self, exchange_id, book):
    self._books[exchange_id] = book

    if self._books[0] is None or self._books[1] is None:
      self._logger.info('book %d is not ready' % (1 - exchange_id))
      return False

    bid_exchange0 = float(self._books[0].get_notional_bids_by_amt(
        self._price_det_notional)[1][-1][0])
    ask_exchange0 = float(self._books[0].get_notional_asks_by_amt(
        self._price_det_notional)[1][-1][0])
    bid_exchange1 = float(self._books[1].get_notional_bids_by_amt(
        self._price_det_notional)[1][-1][0])
    ask_exchange1 = float(self._books[1].get_notional_asks_by_amt(
        self._price_det_notional)[1][-1][0])

    diff_0_1 = bid_exchange0 - ask_exchange1
    diff_1_0 = bid_exchange1 - ask_exchange0
    avg_price = (bid_exchange0 + ask_exchange0 + bid_exchange1 + ask_exchange1) / 4.

    if abs(diff_0_1 / avg_price) > 0.1 or abs(diff_1_0 / avg_price) > 0.1:
      # price is too deviated, 10%, something is wrong
      return False

    self.edge_0_1 = diff_0_1 / avg_price
    self.edge_1_0 = diff_1_0 / avg_price
    self.ask_exchange0 = ask_exchange0
    self.ask_exchange1 = ask_exchange1
    self.bid_exchange0 = bid_exchange0
    self.bid_exchange1 = bid_exchange1
    self.ready = True

    return True


class TrueBookValue(object):
  def __init__(
      self,
      *,
      timestamp,
      bid,  # true
      ask,  # true
      bid_actual,
      ask_actual):
    self.timestamp = timestamp
    self.bid = bid
    self.ask = ask
    self.bid_actual = bid_actual
    self.ask_actual = ask_actual

  @property
  def mid(self):
    return (self.bid + self.ask) / 2.0

  @property
  def mid_actual(self):
    return (self.bid_actual + self.ask_actual) / 2.0

  def get_age(self):
    return get_timestamp() - self.timestamp


class TrueBookCalculator(object):
  def __init__(self, price_det_notional):
    self._price_det_notional = price_det_notional
    self._truebook = None

  @property
  def truebook(self):
    return self._truebook

  def update(self, book):
    timestamp = book.timestamp
    bid = book.get_notional_bids_by_amt(self._price_det_notional)[1][-1][0]
    ask = book.get_notional_asks_by_amt(self._price_det_notional)[1][-1][0]
    self._truebook = TrueBookValue(timestamp=timestamp,
                                   bid=bid,
                                   ask=ask,
                                   bid_actual=book.bid0().price,
                                   ask_actual=book.ask0().price)

    return True


def _calculate_spread(askp_1, bidp_1):
  return (askp_1 - bidp_1) / (bidp_1 + askp_1) * 2


# 0 = target_market
# 1 = ref_market
class UnhedgePriceCalculator(object):
  def __init__(self,
               basis_ma_window,
               tick,
               book_askt_func_0,
               book_bidt_func_0,
               book_askt_func_1,
               book_bidt_func_1,
               ccy_0_1=1):
    self._use_basis = basis_ma_window > 0
    self._basis_ma_window = basis_ma_window
    self._basis_ma_required = 10
    self._tick_passive = FixedTickPrice(tick) if isinstance(tick, float) else tick
    self._book_askt_funcs = {0: book_askt_func_0, 1: book_askt_func_1}
    self._book_bidt_funcs = {0: book_bidt_func_0, 1: book_bidt_func_1}
    self._ask0s = [None, None]
    self._bid0s = [None, None]
    self._askts = [None, None]
    self._bidts = [None, None]
    self._midps = [None, None]

    self._ready = False
    self._last = None
    self._ccy_0_1 = ccy_0_1

    if self._use_basis:
      self._ask_basis_ma = MovingSum(self._basis_ma_window)
      self._bid_basis_ma = MovingSum(self._basis_ma_window)

  @property
  def tick(self):
    return self._tick_passive

  @property
  def ready(self):
    return self._ready

  @property
  def last_value(self):
    return self._last

  def reset(self):
    self._ready = False
    self._last = None

  def update(self, book_0, book_1):
    book_ts = max(book_0.timestamp, book_1.timestamp)
    min_ts = min(book_0.timestamp, book_1.timestamp)

    if book_ts - min_ts > (2 * 10**9):
      self.reset()
      return

    if not (book_0.has_ask() and book_0.has_bid() and book_1.has_ask() and book_1.has_bid()):
      self.reset()
      return

    # Exchange 0
    ask0_0 = book_0.ask0().price * self._ccy_0_1
    bid0_0 = book_0.bid0().price * self._ccy_0_1
    askt_0 = self._book_askt_funcs[0](book_0) * self._ccy_0_1
    bidt_0 = self._book_bidt_funcs[0](book_0) * self._ccy_0_1
    assert askt_0 >= ask0_0 >= bid0_0 >= bidt_0, (askt_0, ask0_0, bid0_0, bidt_0)
    self._ask0s[0] = ask0_0
    self._bid0s[0] = bid0_0
    self._askts[0] = askt_0
    self._bidts[0] = bidt_0
    self._midps[0] = (ask0_0 + bid0_0) / 2

    # Exchange 1
    ask0_1 = book_1.ask0().price
    bid0_1 = book_1.bid0().price
    askt_1 = self._book_askt_funcs[1](book_1)
    bidt_1 = self._book_bidt_funcs[1](book_1)
    assert askt_1 >= ask0_1 >= bid0_1 >= bidt_1, (askt_1, ask0_1, bid0_1, bidt_1)
    self._ask0s[1] = ask0_1
    self._bid0s[1] = bid0_1
    self._askts[1] = askt_1
    self._bidts[1] = bidt_1
    self._midps[1] = (ask0_1 + bid0_1) / 2

    if self._use_basis:
      askt_diff = askt_0 - askt_1
      bidt_diff = bidt_0 - bidt_1

      self._ask_basis_ma.update(book_ts, askt_diff)
      self._bid_basis_ma.update(book_ts, bidt_diff)

      if len(self._ask_basis_ma) < self._basis_ma_required or \
         len(self._bid_basis_ma) < self._basis_ma_required:
        self.reset()
        return

      self._last = FeatureRow(time=book_ts,
                              time_hr=str(to_datetime(book_ts)),
                              ask0_0=ask0_0,
                              askt_0=askt_0,
                              bid0_0=bid0_0,
                              bidt_0=bidt_0,
                              ask0_1=ask0_1,
                              askt_1=askt_1,
                              bid0_1=bid0_1,
                              bidt_1=bidt_1,
                              ask_basis_ma=self._ask_basis_ma.average,
                              bid_basis_ma=self._bid_basis_ma.average)
    else:
      self._last = FeatureRow(time=book_ts,
                              time_hr=str(to_datetime(book_ts)),
                              ask0_0=ask0_0,
                              askt_0=askt_0,
                              bid0_0=bid0_0,
                              bidt_0=bidt_0,
                              ask0_1=ask0_1,
                              askt_1=askt_1,
                              bid0_1=bid0_1,
                              bidt_1=bidt_1,
                              ask_basis_ma=0.,
                              bid_basis_ma=0.)

    self._ready = True

  def adjust_price(self, askp_1, bidp_1):
    askp_1 = self._tick_passive.ceil(askp_1)
    bidp_1 = self._tick_passive.floor(bidp_1)

    ask0_minus_one_1 = self._tick_passive.get_prev_price(self._last.ask0_1)
    bid0_plus_one_1 = self._tick_passive.get_next_price(self._last.bid0_1)

    askp_1 = max(askp_1, ask0_minus_one_1, bid0_plus_one_1)
    bidp_1 = min(bidp_1, ask0_minus_one_1, bid0_plus_one_1)

    # TODO(inkyu): Deprecate return value
    return askp_1, bidp_1

  def calculate_price(self, sell_edge, buy_edge, *, max_spread=None):
    assert self.ready
    assert -0.5 < sell_edge < 0.5
    assert -0.5 < buy_edge < 0.5

    askt_0_proj_1 = self._last.askt_0
    bidt_0_proj_1 = self._last.bidt_0

    if self._use_basis:
      askt_0_proj_1 -= float(self._ask_basis_ma.average)
      bidt_0_proj_1 -= float(self._bid_basis_ma.average)

    askp_1 = askt_0_proj_1 * (1. + sell_edge)
    bidp_1 = bidt_0_proj_1 * (1. - buy_edge)

    return self.adjust_price(askp_1, bidp_1)

  def calculate_price_by_spread(self, spread):
    assert self.ready
    assert 0 < spread < 0.5

    midp_proj_1 = (self._last.askt_0 + self._last.bidt_0) / 2.0

    if self._use_basis:
      midp_proj_1 -= float((self._ask_basis_ma.average + self._bid_basis_ma.average) / 2.0)

    # 0.1bp for buffer
    askp_1 = midp_proj_1 * (1. + spread / 2.0 - 0.00001)
    bidp_1 = midp_proj_1 * (1. - spread / 2.0 + 0.00001)

    cal_spread = _calculate_spread(askp_1, bidp_1)

    assert cal_spread <= spread, (cal_spread, spread)

    return self.adjust_price(askp_1, bidp_1)


class PositionLeanPricer(object):
  def __init__(self, min_pos, max_pos, lot_size, lean_edge_bp, base_edge_bp=0):
    assert max_pos >= min_pos, (max_pos, min_pos)
    self._lot_size = lot_size
    self._base_edge_bp = base_edge_bp
    # set upper limit of lean_edge_bp to 500 to prevent highly skewed price
    # generated by pricer
    self._lean_edge_bp = min(500, lean_edge_bp)
    self._pos_range = max_pos - min_pos
    self._avg_pos = 0.5 * (min_pos + max_pos)
    self._max_lot = 0.5 * self._pos_range / lot_size

  def _get_imbalance_bp(self, pos):
    accum_stack = (pos - self._avg_pos) / self._lot_size
    sign = 0 if accum_stack == 0 else (1 if accum_stack > 0 else -1)
    lean_ratio = abs(accum_stack) / self._max_lot
    return self._lean_edge_bp * min(lean_ratio, 1) * sign

  def update(self, pos, sell_price, buy_price):
    imbalance_bp = self._get_imbalance_bp(pos)
    buy_edge_bp = self._base_edge_bp + imbalance_bp
    sell_edge_bp = self._base_edge_bp - imbalance_bp
    adj_buy_price = None
    adj_sell_price = None
    if buy_price is not None:
      adj_buy_price = buy_price - buy_price * buy_edge_bp * (1e-4)
    if sell_price is not None:
      adj_sell_price = sell_price + sell_price * sell_edge_bp * (1e-4)
    return adj_sell_price, adj_buy_price
